import os
from langchain_openai import ChatOpenAI
from utils.TimeUtils import *
from langchain.prompts import StringPromptTemplate

api_key = os.getenv("OPENAI_API_KEY")
base_url = os.getenv("OPENAI_API_BASE")

def hello_world(abc):
    print("hello world")
    return abc


PROMPT = """
你是一个非常富有经验和天赋的程序员, 现在给你如下的函数名称，你会按照如下格式，输出代码的名称，源代码，中文解释。
函数名称:{function_name}
源代码:
{source_code}
代码解释:
"""

import inspect
def get_source_code(function_name):
    return inspect.getsource(function_name)

class CustmPrompt(StringPromptTemplate):

    def format(self, **kwargs) -> str:
        source_code = get_source_code(kwargs["function_name"])

        print("=============================")
        print(source_code)
        print("=============================")

        prompt = PROMPT.format(
            function_name=kwargs["function_name"].__name__, source_code=source_code
        )
        return prompt

a = CustmPrompt(input_variables=["function_name"])
pm = a.format(function_name=hello_world)

llm = ChatOpenAI(
    model="gpt-3.5-turbo",
    temperature=0,
    api_key=api_key,
    base_url=base_url
)

msg = llm.invoke(pm)
print(msg.content, end="\n")